// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include "function_base_test.h"
#include <pollux/functions/function_registry.h>
#include <pollux/functions/prestosql/registration/registration_functions.h>
#include <pollux/parse/type_resolver.h>

namespace kumo::pollux::functions::test {
    void FunctionBaseTest::SetUpTestCase() {
        parse::registerTypeResolver();
        functions::prestosql::registerAllScalarFunctions();
        memory::MemoryManager::testingSetInstance({});
    }

    // static
    std::vector<const exec::FunctionSignature *> FunctionBaseTest::getSignatures(
        const std::string &functionName,
        const std::string &returnType) {
        const auto allSignatures = get_function_signatures();

        std::vector<const exec::FunctionSignature *> signatures;
        for (const auto &signature: allSignatures.at(functionName)) {
            const auto &typeName = signature->returnType().baseName();
            if (exec::sanitize_name(typeName) == exec::sanitize_name(returnType)) {
                signatures.push_back(signature);
            }
        }
        POLLUX_CHECK(
            !signatures.empty(),
            "No signature found for function {} with return type {}.",
            functionName,
            returnType);
        return signatures;
    }

    // static
    std::unordered_set<std::string> FunctionBaseTest::getSignatureStrings(
        const std::string &functionName) {
        auto allSignatures = get_function_signatures();
        const auto &signatures = allSignatures.at(functionName);

        std::unordered_set<std::string> signatureStrings;
        for (const auto &signature: signatures) {
            signatureStrings.insert(signature->toString());
        }
        return signatureStrings;
    }

    std::pair<VectorPtr, melon::F14FastMap<std::string, exec::ExprStats> >
    FunctionBaseTest::evaluateWithStats(
        const std::string &expression,
        const RowVectorPtr &data,
        const std::optional<SelectivityVector> &rows) {
        auto typedExpr = makeTypedExpr(expression, as_row_type(data->type()));

        std::vector<VectorPtr> results(1);

        exec::ExprSet exprSet({typedExpr}, &execCtx_);
        exec::EvalCtx evalCtx(&execCtx_, &exprSet, data.get());
        if (rows.has_value()) {
            exprSet.eval(*rows, evalCtx, results);
        } else {
            SelectivityVector defaultRows(data->size());
            exprSet.eval(defaultRows, evalCtx, results);
        }

        return {results[0], exprSet.stats()};
    }

    void FunctionBaseTest::testEncodings(
        const core::TypedExprPtr &expr,
        const std::vector<VectorPtr> &inputs,
        const VectorPtr &expected) {
        POLLUX_CHECK(!inputs.empty());

        const auto size = inputs[0]->size();
        POLLUX_CHECK_GE(size, 3);

        auto testDictionary = [&](vector_size_t dictionarySize,
                                  std::function<vector_size_t(vector_size_t)> indexAt,
                                  std::function<bool(vector_size_t)> nullAt =
                                          nullptr) {
            // Wrap each input in its own dictionary.
            std::vector<VectorPtr> encodedInputs;
            encodedInputs.reserve(inputs.size());
            for (const auto &input: inputs) {
                auto indices = make_indices(dictionarySize, indexAt);
                auto nulls =
                        nullAt ? make_nulls(dictionarySize, nullAt) : BufferPtr(nullptr);
                encodedInputs.emplace_back(
                    BaseVector::wrap_in_dictionary(nulls, indices, dictionarySize, input));
            }
            auto encodedRow = make_row_vector(encodedInputs);

            SCOPED_TRACE(fmt::format("Dictionary: {}", encodedRow->toString()));

            auto indices = make_indices(dictionarySize, indexAt);
            auto nulls =
                    nullAt ? make_nulls(dictionarySize, nullAt) : BufferPtr(nullptr);
            auto expectedResult =
                    BaseVector::wrap_in_dictionary(nulls, indices, dictionarySize, expected);
            pollux::test::assertEqualVectors(expectedResult, evaluate(expr, encodedRow));
        };

        auto testConstant = [&](vector_size_t row) {
            std::vector<VectorPtr> constantInputs;
            for (const auto &input: inputs) {
                constantInputs.push_back(BaseVector::wrap_in_constant(100, row, input));
            }
            auto constant_row = make_row_vector(constantInputs);

            SCOPED_TRACE(fmt::format("Constant: {}", constant_row->toString()));

            auto expectedResult = BaseVector::wrap_in_constant(100, row, expected);
            pollux::test::assertEqualVectors(
                expectedResult, evaluate(expr, constant_row));
        };

        SCOPED_TRACE(expr->toString());

        // No extra encoding.
        pollux::test::assertEqualVectors(
            expected, evaluate(expr, make_row_vector(inputs)));

        // Repeat each row twice: 0, 0, 1, 1,... No extra nulls.
        testDictionary(size * 2, [](auto row) { return row / 2; });

        // Select even rows: 0, 2, 4,... No extra nulls.
        testDictionary(size / 2, [](auto row) { return row * 2; });

        // Go over all rows in reverse: N, N-1, N-2,...0. No extra nulls.
        testDictionary(size, [&](auto row) { return size - 1 - row; });

        // Repeat each row twice: 0, 0, 1, 1,... Add some nulls.
        testDictionary(size * 2, [](auto row) { return row / 2; }, null_every(3));

        // Select even rows: 0, 2, 4,... Add some nulls.
        testDictionary(size / 2, [](auto row) { return row * 2; }, null_every(3));

        // Go over all rows in reverse: N, N-1, N-2,...0. Add some nulls.
        testDictionary(size, [&](auto row) { return size - 1 - row; }, null_every(3));

        // Generate constant vectors and verify the results.
        for (auto i = 0; i < size; ++i) {
            testConstant(i);
        }
    }
} // namespace kumo::pollux::functions::test
